import requests
from cozepy import COZE_CN_BASE_URL

class CozeConversation(object):
    def __init__(self):
        self.token = "pat_ODeINGHZ9FvwVWQ7b3mONoWsKKhKjGGyWN8yqBuejQO202EvOrPm49KJ2fCPm1ST"
        self.headers = {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json"
        }
        self.bot_id = "7545456739190456347"     # 智能体id

    def create_conversation(self, question):
        url = "https://api.coze.cn/v1/conversation/create"
        # 请求体
        payload = {
            "bot_id": self.bot_id,
            "name": question
        }
        # 发起请求
        resp = requests.post(url, headers=self.headers, json=payload, timeout=10)
        return resp.json()

    def update_conversation(self, conversation_id, conversation_name):
        url = f"https://api.coze.cn/v1/conversations/{conversation_id}"
        data = {
            "name": conversation_name
        }
        response = requests.put(url=url, headers=self.headers, json=data)
        # print(response.text)
        # print(response.url)
        return response.json()

    def delete_conversation(self, conversation_id):
        url = f"https://api.coze.cn/v1/conversations/{conversation_id}"
        response = requests.delete(url=url, headers=self.headers)
        # print(response.text)
        # print(response.url)
        return response.json()

    def find_all_conversation(self):
        url = "https://api.coze.cn/v1/conversations"
        response = requests.get(url=url, headers=self.headers, params={"bot_id": self.bot_id})
        # print(response.text)
        # print(response.url)
        return response.json()

    def find_message_list(self, conversation_id):
        url = "https://api.coze.cn/v1/conversation/message/list"
        data = {
            "order": "desc",
        }
        response = requests.post(url, headers=self.headers, json=data, params={"conversation_id": conversation_id})
        # print(response.text)
        # print(response.url)
        return response.json()

    def change_message(self, conversation_id, message_id,  content):
        url = 'https://api.coze.cn/v1/conversation/message/modify'
        body = {
            "content": content,
            "content_type": "text"
        }
        print(conversation_id, message_id, content)
        response = requests.post(url, headers=self.headers, json=body, params={"conversation_id": str(conversation_id), "message_id": str(message_id)})
        return response.json()
    # 非流式输出
    # def have_conversation(self, user_id, conversation_id, question):
    #     from cozepy import COZE_CN_BASE_URL
    #     messages = ""
    #     coze_api_token = self.token
    #     coze_api_base = COZE_CN_BASE_URL
    #     from cozepy import Coze, TokenAuth, Message, ChatStatus, MessageContentType, ChatEventType  # noqa
    #     coze = Coze(auth=TokenAuth(token=coze_api_token), base_url=coze_api_base)
    #     for event in coze.chat.stream(
    #             bot_id=self.bot_id,
    #             user_id=str(user_id),
    #             conversation_id=conversation_id,
    #             additional_messages=[
    #                 Message.build_user_question_text(question),
    #             ],
    #     ):
    #         if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
    #             # print(event.message.content, end="", flush=True)
    #             messages += event.message.content
    #
    #         if event.event == ChatEventType.CONVERSATION_CHAT_COMPLETED:
    #             print()
    #             # print("token usage:", event.chat.usage.token_count)
    #     return messages


    # 流式输出
    def have_conversation(self, user_id, conversation_id, question, clear_the_message=None):
        # messages = ""
        coze_api_token = self.token
        coze_api_base = COZE_CN_BASE_URL
        from cozepy import Coze, TokenAuth, Message, ChatStatus, MessageContentType, ChatEventType  # noqa
        coze = Coze(auth=TokenAuth(token=coze_api_token), base_url=coze_api_base)
        def chat_stream():
            for event in coze.chat.stream(
                    bot_id=self.bot_id,
                    user_id=str(user_id),
                    conversation_id=conversation_id,
                    additional_messages=[
                        Message.build_user_question_text(str(question)),
                    ],
            ):
                if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
                    # print(event.message.content, end="", flush=True)
                    yield f"data: {event.message.content}\n\n\n"
                elif event.event == ChatEventType.CONVERSATION_CHAT_COMPLETED:
                    if clear_the_message:
                        message_id = self.find_message_list(conversation_id).get("data")[1].get("id")  # 获取最新的user消息id
                        print(self.change_message(conversation_id, message_id, " "))  # 将上面获取的消息改为空
                    # 可选：结束标记
                    yield f"data: [DONE]\n\n"
        return chat_stream



if __name__ == '__main__':
    # coze = CozeConversation()
    # coze.have_conversation(user_id=1, conversation_id=7541241130983571483, question="你好，可以给我推荐一下电脑吗")
    coze = CozeConversation()
    print(coze.change_message(conversation_id=7542552795548581938, message_id=7542555642406633472, content="hh"))